from kuai.sim import *
from kuai.mol.io import read_molfile
from kuai.sim.dff2kuai import read_ppffile
from kuai.unit import default_units
from numpy import zeros, double, sum, random, reshape, matrix
from scipy import optimize 
from time import clock

class EnergyFunction:
    def __init__(self, model, funcs, job):
        self.model = model
        self.funcs = funcs
        self.job = job

    def __call__(self, x):
        self.model.coords = x
        e = self.job.getE(self.model, self.funcs)
        result = sum(e)
        return result

        
class GradientFunction:
    def __init__(self, model, funcs, job):
        self.model = model
        self.funcs = funcs
        self.job = job
        
    def __call__(self, x):
        self.model.coords = x
        _, g = self.job.getEG(self.model, self.funcs)
        return g
        
class HessianFunction:
    def __init__(self, model, funcs, job):
        self.model = model
        self.funcs = funcs
        self.job = job
        
    def __call__(self, x):
        self.model.coords = x
        _, _, h = self.job.getEGH(self.model, self.funcs)
        result = reshape(h, (len(self.model.coords), len(self.model.coords)))
        return result


def test(model, index, startX):
    funcs = make_efunc(model, index)
    job = SimulationJob()
    efunc = EnergyFunction(model, funcs, job)
    gfunc = GradientFunction(model, funcs, job)
    hfunc = HessianFunction(model, funcs, job)
    
    fe = default_units.format(1, 'kcal/mol')[0] # Factor of Energy Unit
    fg = default_units.format(1, 'kcal/mol A')[0] # Factor of Energy Unit
    
    e0 = job.getE(model, funcs) * fe
    print e0
    e0 = efunc(startX) * fe 
    print "\n\n\n\n"
    print "=" * 72
    print "Init energy = ", e0, 'kcal/mol'

    """
    print "\n\nDownhill Simplex Algorithm"
    t1 = clock()
    result = optimize.fmin(efunc, startX, full_output=True)
    print "Cost", clock()-t1
    print "Final Energy = ", result[1] * fe, 'kcal/mol'
    
    print "\n\nPowell Algorithm"
    t1 = clock()
    result = optimize.fmin_powell(efunc, startX, full_output=True)
    print "Cost", clock()-t1
    print "Final Energy = ", result[1] * fe, 'kcal/mol'

    print "\n\n Conjugate Gradient Algorithm"
    t1 = clock()
    result = optimize.fmin_cg(efunc, startX, gfunc, epsilon=1e-6, full_output=True)
    print "Cost", clock()-t1
    print "Final Energy = ", result[1] * fe, 'kcal/mol'
    
    print "\n\n Conjugate Gradient Algorithm"
    t1 = clock()
    result = optimize.fmin_cg(efunc, startX, epsilon=1e-6, full_output=True)
    print "Cost", clock()-t1
    print "Final Energy = ", result[1] * fe, 'kcal/mol'
    
    print "\n\n BFGS Algorithm"
    t1 = clock()
    result = optimize.fmin_bfgs(efunc, startX, gfunc, epsilon=1e-6, full_output=True)
    print "Cost", clock()-t1
    print "Final Energy = ", result[1] * fe, 'kcal/mol'

    print "\n\n BFGS Algorithm"
    t1 = clock()
    result = optimize.fmin_bfgs(efunc, startX, epsilon=1e-6, full_output=True)
    print "Cost", clock()-t1
    print "Final Energy = ", result[1] * fe, 'kcal/mol'
    """

    
    print "\n\n Newtown Conjugate Gradient Algorithm"
    t1 = clock()
    result = optimize.fmin_ncg(efunc, startX, gfunc, fhess=hfunc, epsilon=1e-5, full_output=True)
    print "Cost", clock()-t1
    print "Final Energy = ", result[1] * fe, 'kcal/mol'
    
    """
    print "\n\n Simulated Annealing"
    t1 = clock()
    result = optimize.anneal(efunc, startX, full_output=True)
    print "Cost", clock()-t1
    print "Final Energy = ", result[1] * fe, 'kcal/mol'
    """

if __name__ == '__main__':
    import sys
    mol = read_molfile(sys.argv[1])
    index, parameters = read_ppffile(sys.argv[2])
    model = setup_model(mol, index, parameters)
    
    startX = model.coords
    test(model, index, startX)
    
    for i in range(10):
        random.seed(i)
        startX =random.rand(len(startX)) * 10
        test(model, index, startX)
        
